//===------ AccessEnforcementDom.cpp - dominated access removal opt -------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
///
/// This function pass removes dominated accesses in two ways:
///
/// 1) Remove accesses dominated by an existing access
///
/// General case:
/// begin_access A (may or may not have no_nested_conflict)
/// load/store
/// end_access
/// ...
/// begin_access A [no_nested_conflict] // dominated by the first access
/// load/store
/// end_access A
///
/// The second access scope does not need to be emitted.
///
/// 2) Add a new dominating accesses to loop's preheader
///
/// General case:
/// <loop preheader>
/// A = ref_element_addr
/// <loop>
/// begin_access A [dynamic] [no_nested_conflict]
///
/// Adding an empty begin_access A in the preheader would allow us to
/// turn the loop's access to [static]
///
/// Warning: This optimization requires that all points within this function
/// that begin an access can be identified. Failure to recognize the beginning
/// of an access scope could weaken dynamic enforcement.
///
/// FIXME: This pass currently only runs in the last-chance pipeline, with a
/// guarantee that no access marker removal is done after it. This happens to
/// work but is dangerous and violates PIL semantics. We should instead add a
/// flag for accesses to give them the semantics that they may guard memory
/// operations other than those enclosed by the access scope.
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "access-enforcement-dom"

#include "polarphp/pil/lang/DebugUtils.h"
#include "polarphp/pil/lang/MemAccessUtils.h"
#include "polarphp/pil/lang/PILBuilder.h"
#include "polarphp/pil/lang/PILFunction.h"
#include "polarphp/pil/optimizer/analysis/DominanceAnalysis.h"
#include "polarphp/pil/optimizer/analysis/LoopAnalysis.h"
#include "polarphp/pil/optimizer/passmgr/Transforms.h"
#include "polarphp/pil/optimizer/utils/InstOptUtils.h"
#include "llvm/ADT/DepthFirstIterator.h"

using namespace polar;

// =============================================================================
// DominatedAccessAnalysis

namespace polar {

/// Information about each dynamic access with valid storage.
///
/// This is a pass-specific subclass of AccessedStorage with identical layout.
/// An instance is created for each BeginAccess in the current function. In
/// additional to identifying the access' storage location, it associates that
/// access with pass-specific data in reserved bits. The reserved bits do not
/// participate in equality or hash lookup.
class DomAccessedStorage : public AccessedStorage {
public:
   DomAccessedStorage() {}

   explicit DomAccessedStorage(const AccessedStorage &storage)
      : AccessedStorage(storage) {
      Bits.DomAccessedStorage.isInner = false;
      Bits.DomAccessedStorage.containsRead = false;
   }

   // Is this dynamic, identifiable access scope potentially contained within any
   // kind of outer scope. The outer scope may be static and/or have unidentified
   // storage. For example:
   //
   // A1: [static] [read]
   // A2: [dynamic] [read]
   // A3: [dynamic] [modify]
   //
   // A2 cannot be promoted to modify if its scope overlaps with A1.
   bool isInner() const { return Bits.DomAccessedStorage.isInner; }

   void setIsInner() { Bits.DomAccessedStorage.isInner = true; }

   /// Is this dynamic, identifiable access scope a [read], and does it
   /// potentially contain another non-distinct [read] access of any kind?
   bool containsRead() const { return Bits.DomAccessedStorage.containsRead; }

   void setContainsRead() { Bits.DomAccessedStorage.containsRead = true; }

   void dump() const {
      AccessedStorage::dump();
      llvm::dbgs() << "<" << (isInner() ? "" : "inner")
                   << (containsRead() ? "" : "containsRead") << ">\n";
   }
};
} // namespace polar

namespace {
// An analysis that maps each valid dynamic BeginAccess to
// DomAccessedStorage. Performs a trivial data flow analysis to populate the map
// with information about nested accesses. Data flow is needed to track open
// access scopes, but only flow-insensitive information is recorded in the
// result.
//
// Note that all access scopes are tracked during data flow, but only valid
// dynamic are mapped to results.
//
// TODO: Separate this into a shared analysis and factor it with
// AccessEnforcementOpts. The optimization that merges accesses would also
// benefit.
//
// TODO: This could be made more precise by querying AccessedStorageAnalysis.
class DominatedAccessAnalysis {
public:
   // The result records information for all dynamic accesses in this
   // function. If an UnpairedAccess exists, then the result will be
   // consevatively empty.
   struct Result {
      llvm::SmallDenseMap<BeginAccessInst *, DomAccessedStorage, 32> accessMap;
   };

private:
   // The data flow state for each block. It is only valid between the time at
   // least one predecesor block has been finished and before its own block
   // has been finished.
   //
   // The isBottom flag allows the analysis to avoid quadratic behavior where
   // each open access must be updated at each conservatively handled
   // instruction. Instead, the accesses are only conservatively updated once
   // until the next scope is entered. So the complexity is (#OpenScopes)^2
   // instead of (#OpenScopes)*(#Applies).
   struct BBState {
      using DenseAccessSet = llvm::SmallDenseSet<BeginAccessInst *, 4>;
      using DenseCoroutineSet = llvm::SmallDenseSet<BeginApplyInst *, 4>;

      DenseAccessSet inScopeAccesses;
      DenseCoroutineSet inScopeCoroutines;
      bool isBottom = false;
   };
   using BlockStateMap = llvm::DenseMap<PILBasicBlock *, BBState>;

   PostOrderFunctionInfo *PO;

   Result result; // Flow-insensitive analysis result.

   BlockStateMap blockStateMap; // Data flow state.

public:
   DominatedAccessAnalysis(PostOrderFunctionInfo *PO) : PO(PO) {}

   Result analyze() &&;

protected:
   void setBottom(BBState &state);
   void analyzeAccess(BeginAccessInst *BAI, BBState &state);
};
} // namespace

// Set information for all in-scope accesses to the worst case "bottom".
// The isInner flag is not affected because it only applies to scopes inside the
// currently open scopes, not the currently open scopes themselves.
void DominatedAccessAnalysis::setBottom(BBState &state) {
   if (state.isBottom)
      return;

   // Unordered iteration over the in-access scopes.
   llvm::for_each(state.inScopeAccesses, [this](BeginAccessInst *BAI) {
      if (auto &domStorage = result.accessMap[BAI])
         domStorage.setContainsRead();
   });
}

// Perform the analysis and return the Result, relinquishing internal state.
DominatedAccessAnalysis::Result DominatedAccessAnalysis::analyze() && {
   // A single RPO traversal is sufficient to visit access scopes in order. The
   // set of open accesses entering a block cannot grow as a result of loops, and
   // within an open scope the results are flow-insensitive.
   for (auto *BB : PO->getReversePostOrder()) {
      BBState state = blockStateMap[BB];
      for (auto &I : *BB) {
         if (auto *BAI = dyn_cast<BeginAccessInst>(&I)) {
            analyzeAccess(BAI, state);
            continue;
         }
         if (auto *EAI = dyn_cast<EndAccessInst>(&I)) {
            bool erased = state.inScopeAccesses.erase(EAI->getBeginAccess());
            (void)erased;
            assert(erased);
            continue;
         }
         // Check for BeginApply before checking FullApplySite below.
         if (auto *beginApply = dyn_cast<BeginApplyInst>(&I)) {
            auto iterAndInserted = state.inScopeCoroutines.insert(beginApply);
            (void)iterAndInserted;
            assert(iterAndInserted.second);
            continue;
         }
         if (auto *endApply = dyn_cast<EndApplyInst>(&I)) {
            bool erased = state.inScopeCoroutines.erase(endApply->getBeginApply());
            (void)erased;
            assert(erased);
            continue;
         }
         if (FullApplySite::isa(&I)) {
            setBottom(state);
            continue;
         }
         if (isa<BeginUnpairedAccessInst>(&I)) {
            // Unpaired accesses could be tracked, but are ignored because they are
            // mostly irrelevant and hard to test. Completely bail on this function.
            result.accessMap.clear();
            return std::move(result);
         }
      }
      auto successors = BB->getTerminator()->getSuccessors();
      unsigned numSucc = successors.size();
      for (unsigned succIdx : indices(successors)) {
         PILBasicBlock *succBB = successors[succIdx].getBB();
         if (succBB == BB)
            continue;

         if (succIdx != numSucc - 1)
            blockStateMap.try_emplace(succBB, state);
         else
            // Move the state into the last successor to avoid copying sets.
            blockStateMap.try_emplace(succBB, std::move(state));
      }
   }
   return std::move(result);
}

// The data flow transfer function for BeginAccess. Creates the
// DomAccessedStorage for this access and inserts it in the result.
void DominatedAccessAnalysis::analyzeAccess(BeginAccessInst *BAI,
                                            BBState &state) {
   DomAccessedStorage domStorage;
   // Only track dynamic access in the result. Static accesses still need to be
   // tracked by data flow, but they can't be optimized as "dominating".
   if (BAI->getEnforcement() == PILAccessEnforcement::Dynamic) {
      AccessedStorage storage = findAccessedStorageNonNested(BAI->getSource());
      // Copy the AccessStorage into DomAccessedStorage. All pass-specific bits
      // are initialized to zero.
      domStorage = DomAccessedStorage(storage);
   }
   // Continue to handle both untracked access and invalid domStorage
   // conservatively below...

   // unordered set iteration...
   llvm::for_each(state.inScopeAccesses, [&](BeginAccessInst *outerBegin) {
      auto &outerInfo = result.accessMap[outerBegin];
      // If the current access is mapped, set its isInner flag.
      if (domStorage && !domStorage.isInner()) {
         if (domStorage.isDistinctFrom(outerInfo))
            return;

         domStorage.setIsInner();
      }
      // The results for tracked in-scope accesses still need to be
      // updated even if the current access is not be tracked.
      if (outerInfo && BAI->getAccessKind() == PILAccessKind::Read
          && outerBegin->getAccessKind() == PILAccessKind::Read) {
         outerInfo.setContainsRead();
      }
   });
   // Track this access even if it is invalid or unmapped.
   {
      auto iterAndInserted = state.inScopeAccesses.insert(BAI);
      (void)iterAndInserted;
      assert(iterAndInserted.second);
   }
   // Update the results if this access will be mapped.
   if (!domStorage)
      return;

   // Set the current accesss isInner flag if it's inside a coroutine scope.
   if (!state.inScopeCoroutines.empty())
      domStorage.setIsInner();

   // Map the current access.
   {
      auto iterAndInserted = result.accessMap.try_emplace(BAI, domStorage);
      (void)iterAndInserted;
      assert(iterAndInserted.second);
   }
   state.isBottom = false;
}

// =============================================================================
// DominatedAccessRemoval optimization.

namespace {
using DomTreeNode = llvm::DomTreeNodeBase<PILBasicBlock>;

// Visit the dominator tree top down, tracking the current set of dominating
// dynamic accesses. Dominated dynamic accesses with identical storage are
// marked static during traversal. If a dynamic access inside a loop has no
// dominating access, insert a new access in the preheader.
class DominatedAccessRemoval {
   // Record the first access of a given storage location and the dominator node
   // in which the access occurred.
   struct DominatingAccess {
      BeginAccessInst *beginAccess;
      DomTreeNode *domNode;

      DominatingAccess(BeginAccessInst *beginAccess, DomTreeNode *domNode)
         : beginAccess(beginAccess), domNode(domNode) {}
   };
   using StorageToDomMap = llvm::DenseMap<AccessedStorage, DominatingAccess>;

   PILFunction &func;
   DominanceInfo *domInfo;
   PILLoopInfo *loopInfo;
   DominatedAccessAnalysis::Result &DAA;

   // Hash map from each storage location to the dominating access.
   StorageToDomMap storageToDomMap;

   DomTreeNode *currDomNode = nullptr;

   bool hasChanged = false;

public:
   DominatedAccessRemoval(PILFunction &func, DominanceInfo *domInfo,
                          PILLoopInfo *loopInfo,
                          DominatedAccessAnalysis::Result &DAA)
      : func(func), domInfo(domInfo), loopInfo(loopInfo), DAA(DAA) {}

   bool optimize();

protected:
   void visitBeginAccess(BeginAccessInst *BAI);
   bool checkDominatedAccess(BeginAccessInst *BAI,
                             DomAccessedStorage currDomStorage);
   bool optimizeDominatedAccess(BeginAccessInst *currBegin,
                                DomAccessedStorage currDomStorage,
                                const DominatingAccess &domAccess);
   void tryInsertLoopPreheaderAccess(BeginAccessInst *BAI,
                                     DomAccessedStorage currDomStorage);
};
} // namespace

// Optimize the current function, and return true if any optimization was
// performed.
bool DominatedAccessRemoval::optimize() {
   DomTreeNode *entryNode = domInfo->getNode(func.getEntryBlock());
   for (DomTreeNode *domNode : llvm::depth_first(entryNode)) {
      currDomNode = domNode;

      // Optimize dominated accesses in this block.
      for (auto &instr : *domNode->getBlock()) {
         if (auto *BAI = dyn_cast<BeginAccessInst>(&instr))
            visitBeginAccess(BAI);
      }
   }
   return hasChanged;
}

// Visit a BeginAccessInst once-and-only-once in domtree order.
// Attempt to find a dominating access with identical storage.
// If that fails, attempt to insert a new dominating access in the preheader.
void DominatedAccessRemoval::visitBeginAccess(BeginAccessInst *BAI) {
   if (BAI->getEnforcement() != PILAccessEnforcement::Dynamic)
      return;

   DomAccessedStorage currDomStorage = DAA.accessMap.lookup(BAI);
   if (!currDomStorage)
      return;

   // Only track "identifiable" storage.
   if (currDomStorage.isUniquelyIdentifiedOrClass()) {
      if (checkDominatedAccess(BAI, currDomStorage))
         return;
   }
   tryInsertLoopPreheaderAccess(BAI, currDomStorage);
}

// Track this identifiable dynamic access in storageToDomMap, and optimize it if
// possible. Return true if the optimization suceeds.
bool DominatedAccessRemoval::checkDominatedAccess(
   BeginAccessInst *BAI, DomAccessedStorage currDomStorage) {
   // Attempt to add this access to storageToDomMap using its base storage
   // location as the key.
   //
   // Cast this DomAccessedStorage back to a plain storage location. The
   // pass-specific bits will be ignored, but reset them anyway for sanity.
   AccessedStorage storage = static_cast<AccessedStorage>(currDomStorage);
   storage.resetSubclassData();
   auto iterAndInserted =
      storageToDomMap.try_emplace(storage, DominatingAccess(BAI, currDomNode));
   if (iterAndInserted.second)
      return false;

   // An access has already been recorded for this storage.
   // If the previous domNode does not dominate currDomNode, then replace it.
   DominatingAccess &domAccess = iterAndInserted.first->second;
   if (!domInfo->dominates(domAccess.domNode, currDomNode)) {
      domAccess = DominatingAccess(BAI, currDomNode);
      return false;
   }
   // The previously mapped access still dominates this block, so the current
   // access can potentially be optimized.
   return optimizeDominatedAccess(BAI, currDomStorage, domAccess);
}

// If possible, optimize the current access by converting it to [static]. Return
// true if the optimization succeeds.
//
// This function is not allowed to add or erase instructions, only change
// instruction flags.
//
// The four required conditions for converting this access to static are:
//
// 1. A closed dominating access has identical storage.
//
// The caller looked up this access' storage in storageToDomMap and checked
// that the previously seen access dominates this block. As long as this access'
// isInner flag from DominatedAccessAnalysis is not set, the dominating access
// must be closed.
//
// 2. There is no open (overlapping) access with nondistinct storage (that isn't
// also open at the dominating access).
//
// The isInner flag from DominatedAccessAnalysis indicates no enclosing
// nondistinct scopes within this function. Any outer scope would enclose the
// whole function, thereby also enclosing the dominating scope.
//
// 3. This access has no_nested_conflict.
//
// This is a direct check on this BeginAccessInst flag.
//
// 4. The current and dominating access kinds are compatible:
//
// read -> read: OK
//
// modify -> read: OK
//
// modify -> modify: OK
//
// read -> modify: Requires promoting the dominating access to a modify. This
// can be done as long as the dominating access does not contain another
// non-distinct read and isn't contained by another non-distinct read. The
// containsRead and isInner flags from in DominatedAccessAnalysis answer this
// conservatively.
//
// Note: Promoting an earlier access to a modify could cause a program to
// trap when optimized even if the unoptimized program does not trap; the
// original modify access may be on an unreachable code path. This is acceptable
// because:
//
// (a) in theory, exclusivity violations do not need to be executed to be
// considered program violations. Promoting the access does not introduce any
// new conflict where one didn't already exist statically. Catching these
// violations at runtime is only an implementation compromise, and the more true
// violations are caught, the better.
//
// (b) in practice, this situation is so exceedingly unlikely that it won't
// cause any pervasive usability problem where programs have stronger
// enforcement only when optimized.
bool DominatedAccessRemoval::optimizeDominatedAccess(
   BeginAccessInst *BAI, DomAccessedStorage currAccessInfo,
   const DominatingAccess &domAccess) {
   // 1. and 2. If any nondistinct scopes are open, it must remain dynamic.
   if (currAccessInfo.isInner())
      return false;

   // 3. If BAI may have a nested conflict, it must remain dynamic.
   if (!BAI->hasNoNestedConflict())
      return false;

   // 4. Promoting a read to a modify is only safe with no nested reads.
   if (domAccess.beginAccess->getAccessKind() == PILAccessKind::Read
       && BAI->getAccessKind() == PILAccessKind::Modify) {
      DomAccessedStorage domStorage = DAA.accessMap[domAccess.beginAccess];
      if (domStorage.containsRead() || domStorage.isInner())
         return false;

      LLVM_DEBUG(llvm::dbgs()
                    << "Promoting to modify: " << *domAccess.beginAccess << "\n");
      domAccess.beginAccess->setAccessKind(PILAccessKind::Modify);
   }
   LLVM_DEBUG(llvm::dbgs() << "Setting static enforcement: " << *BAI << "\n");
   LLVM_DEBUG(llvm::dbgs() << "Dominated by: " << *domAccess.beginAccess
                           << "\n");
   BAI->setEnforcement(PILAccessEnforcement::Static);

   hasChanged = true;
   return true;
}

// Attempt to insert a new access in the loop preheader. If successful, insert
// the new access in DominatedAccessAnalysis so it can be used to dominate other
// accesses. Also convert the current access to static and update the current
// storageToDomMap since the access may already have been recorded (when it was
// still dynamic).
//
// This function cannot add or remove instructions in the current block, but
// may add instructions to the current loop's preheader.
//
// The required conditions for inserting a new dominating access are:
//
// 1. The new preheader access is not enclosed in another scope that doesn't
// also enclose the current scope.
//
// This is inferred from the loop structure; any scope that encloses the
// preheader must also enclose the entire loop.
//
// 2. The current access is not enclosed in another scope that doesn't also
// enclose the preheader.
//
// As before, it is sufficient to check this access' isInner flags in
// DominatedAccessAnalysis; if this access isn't enclosed by any scope within
// the function, then it can't be enclosed within a scope inside the loop.
//
// 3. The current header has no nested conflict within its scope.
//
// 4. The access' source operand is available in the loop preheader.
void DominatedAccessRemoval::tryInsertLoopPreheaderAccess(
   BeginAccessInst *BAI, DomAccessedStorage currAccessInfo) {
   // 2. the current access may be enclosed.
   if (currAccessInfo.isInner())
      return;

   // 3. the current access must be instantaneous.
   if (!BAI->hasNoNestedConflict())
      return;

   PILLoop *currLoop = loopInfo->getLoopFor(BAI->getParent());
   if (!currLoop)
      return;
   PILBasicBlock *preheader = currLoop->getLoopPreheader();
   if (!preheader)
      return;

   // 4. The source operand must be available in the preheader.
   auto sourceOperand = BAI->getOperand();
   auto *sourceBB = sourceOperand->getParentBlock();
   if (!domInfo->dominates(sourceBB, preheader))
      return;

   // Insert a new access scope immediately before the
   // preheader's terminator.
   TermInst *preheaderTerm = preheader->getTerminator();
   PILBuilderWithScope scopeBuilder(preheaderTerm);
   BeginAccessInst *newBegin = scopeBuilder.createBeginAccess(
      preheaderTerm->getLoc(), sourceOperand, BAI->getAccessKind(),
      PILAccessEnforcement::Dynamic, true /*no nested conflict*/,
      BAI->isFromBuiltin());
   scopeBuilder.createEndAccess(preheaderTerm->getLoc(), newBegin, false);
   LLVM_DEBUG(llvm::dbgs() << "Created loop preheader access: " << *newBegin
                           << "\n"
                           << "dominating: " << *BAI << "\n");
   BAI->setEnforcement(PILAccessEnforcement::Static);

   hasChanged = true;

   // Insert the new dominating instruction in both DominatedAccessAnalysis and
   // storageToDomMap if it has uniquely identifiable storage.
   if (!currAccessInfo.isUniquelyIdentifiedOrClass())
      return;

   AccessedStorage storage = static_cast<AccessedStorage>(currAccessInfo);
   storage.resetSubclassData();

   // Create a DomAccessedStorage for the new access with no flags set.
   DAA.accessMap.try_emplace(newBegin, DomAccessedStorage(storage));

   // Track the new access as long as no other accesses from the same storage are
   // already tracked. This also necessarily replaces the current access, which
   // was just made static.
   DominatingAccess newDomAccess(newBegin, domInfo->getNode(preheader));
   auto iterAndInserted = storageToDomMap.try_emplace(storage, newDomAccess);
   if (!iterAndInserted.second) {
      DominatingAccess &curDomAccess = iterAndInserted.first->second;
      if (curDomAccess.beginAccess == BAI)
         curDomAccess = newDomAccess;
   }
}

namespace {
struct AccessEnforcementDom : public PILFunctionTransform {
   void run() override;
};
} // namespace

void AccessEnforcementDom::run() {
   PILFunction *func = getFunction();
   if (func->empty())
      return;

   PostOrderFunctionInfo *PO = getAnalysis<PostOrderAnalysis>()->get(func);
   auto DAA = DominatedAccessAnalysis(PO).analyze();

   DominanceAnalysis *domAnalysis = getAnalysis<DominanceAnalysis>();
   DominanceInfo *domInfo = domAnalysis->get(func);
   PILLoopAnalysis *loopAnalysis = PM->getAnalysis<PILLoopAnalysis>();
   PILLoopInfo *loopInfo = loopAnalysis->get(func);

   DominatedAccessRemoval eliminationPass(*func, domInfo, loopInfo, DAA);
   if (eliminationPass.optimize())
      invalidateAnalysis(PILAnalysis::InvalidationKind::Instructions);
}

PILTransform *polar::createAccessEnforcementDom() {
   return new AccessEnforcementDom();
}
